#!/usr/bin/env python3
"""
数据集分割脚本
将指定目录中的数据按 9:1 比例分割为训练集和验证集
"""

import os
import random
import shutil
import argparse
from pathlib import Path

def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='数据集分割脚本')
    parser.add_argument('--source_dir',
                        type=str,
                        default="/home/dd/working/data/dexin/dexin5/dexined_hy",
                        help='源数据目录路径')
    parser.add_argument('--output_dir',
                        type=str,
                        default="/home/dd/working/data/dexin/dexin5/dexined_hy_split",
                        help='输出目录路径')
    parser.add_argument('--train_ratio',
                        type=float,
                        default=0.9,
                        help='训练集比例 (0.0-1.0)')
    parser.add_argument('--random_seed',
                        type=int,
                        default=42,
                        help='随机种子，确保可重复性')
    
    return parser.parse_args()

def split_dataset_9_1(source_dir, output_dir, train_ratio=0.9, random_seed=42):
    """
    将数据集按指定比例分割为训练集和验证集
    
    Args:
        source_dir: 源数据目录
        output_dir: 输出目录 
        train_ratio: 训练集比例，默认0.9
        random_seed: 随机种子，默认42
    """
    source_path = Path(source_dir)
    output_path = Path(output_dir)
    
    # 检查源目录结构
    images_dir = source_path / "train" / "images"
    gt_dir = source_path / "train" / "gt"
    train_list_file = source_path / "train_pair.lst"
    
    if not all([images_dir.exists(), gt_dir.exists(), train_list_file.exists()]):
        raise ValueError(f"源目录结构不正确: {source_dir}")
    
    # 读取原始训练列表
    with open(train_list_file, 'r') as f:
        lines = [line.strip() for line in f if line.strip()]
    
    print(f"总共找到 {len(lines)} 个样本对")
    
    # 提取文件名对（图像和对应的gt文件）
    sample_pairs = []
    for line in lines:
        parts = line.split()
        if len(parts) >= 2:
            img_path = parts[0]
            gt_path = parts[1]
            # 从路径中提取文件名
            img_name = os.path.basename(img_path)
            gt_name = os.path.basename(gt_path)
            sample_pairs.append((img_name, gt_name))
    
    print(f"成功解析 {len(sample_pairs)} 个文件名对")
    
    # 随机打乱并分割
    random.seed(random_seed)  # 设置随机种子确保可重复性
    random.shuffle(sample_pairs)
    
    train_count = int(len(sample_pairs) * train_ratio)
    train_pairs = sample_pairs[:train_count]
    val_pairs = sample_pairs[train_count:]
    
    print(f"训练集: {len(train_pairs)} 个样本")
    print(f"验证集: {len(val_pairs)} 个样本")
    
    # 创建输出目录结构
    output_path.mkdir(parents=True, exist_ok=True)
    
    # 训练集目录
    train_images_dir = output_path / "train" / "images"
    train_gt_dir = output_path / "train" / "gt"
    train_images_dir.mkdir(parents=True, exist_ok=True)
    train_gt_dir.mkdir(parents=True, exist_ok=True)
    
    # 验证集目录
    val_images_dir = output_path / "val" / "images"
    val_gt_dir = output_path / "val" / "gt"
    val_images_dir.mkdir(parents=True, exist_ok=True)
    val_gt_dir.mkdir(parents=True, exist_ok=True)
    
    # 复制训练集文件
    train_list_pairs = []
    for img_name, gt_name in train_pairs:
        src_img = images_dir / img_name
        src_gt = gt_dir / gt_name
        
        if src_img.exists() and src_gt.exists():
            dst_img = train_images_dir / img_name
            dst_gt = train_gt_dir / gt_name
            
            shutil.copy2(src_img, dst_img)
            shutil.copy2(src_gt, dst_gt)
            
            # 记录训练对
            train_list_pairs.append(f"{dst_img} {dst_gt}")
        else:
            print(f"警告: 文件不存在，跳过 {img_name} -> {gt_name}")
    
    # 复制验证集文件
    val_list_pairs = []
    for img_name, gt_name in val_pairs:
        src_img = images_dir / img_name
        src_gt = gt_dir / gt_name
        
        if src_img.exists() and src_gt.exists():
            dst_img = val_images_dir / img_name
            dst_gt = val_gt_dir / gt_name
            
            shutil.copy2(src_img, dst_img)
            shutil.copy2(src_gt, dst_gt)
            
            # 记录验证对
            val_list_pairs.append(f"{dst_img} {dst_gt}")
        else:
            print(f"警告: 文件不存在，跳过 {img_name} -> {gt_name}")
    
    # 生成训练集列表文件
    train_list_path = output_path / "train_pair.lst"
    with open(train_list_path, 'w') as f:
        for pair in train_list_pairs:
            f.write(pair + '\n')
    
    # 生成验证集列表文件
    val_list_path = output_path / "val_pair.lst"
    with open(val_list_path, 'w') as f:
        for pair in val_list_pairs:
            f.write(pair + '\n')
    
    print(f"\n数据分割完成!")
    print(f"输出目录: {output_path}")
    print(f"训练集列表: {train_list_path} ({len(train_list_pairs)} 对)")
    print(f"验证集列表: {val_list_path} ({len(val_list_pairs)} 对)")
    
    return str(train_list_path), str(val_list_path)

if __name__ == "__main__":
    args = parse_args()
    
    print("开始数据集分割...")
    print(f"源目录: {args.source_dir}")
    print(f"输出目录: {args.output_dir}")
    print(f"训练集比例: {args.train_ratio}")
    print(f"随机种子: {args.random_seed}")
    print("-" * 50)
    
    train_list, val_list = split_dataset_9_1(
        args.source_dir, 
        args.output_dir, 
        args.train_ratio, 
        args.random_seed
    )
    print("\n数据集分割完成!")
    print(f"训练集列表: {train_list}")
    print(f"验证集列表: {val_list}")